import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential, ReLU
import torch.nn.functional as F


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.seq1 = Sequential(
            Conv2d(1, 16, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Conv2d(16, 32, 5, padding=2),
            ReLU(),
            MaxPool2d(2),
            Flatten(),
            Linear(1568, 256),
            ReLU(),
            Linear(256, 64),
            ReLU(),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.seq1(x)
        x = F.softmax(x, dim=1)

        return x


# model = MyModel()
# ins = torch.ones((1, 1, 28, 28))
# print(model(ins).shape)
